import os
import sys
import subprocess
import datetime
from colorama import Fore, Style
from cmd import Cmd
from datetime import datetime

class AiShell(Cmd):
    def __init__(self):
        super().__init__()
        self.prompt = ">>> "
        # 兼容 Windows 和 Linux/Unix 的历史记录文件路径
        self.history_file = os.path.join(os.path.expanduser("~"), ".aita", "shell_history")
        self.history = []
        self._init_history()
        self._print_welcome_message()
        self.agent = None
        self.code_agent = None
        self.current_mode = "normal"  # normal, code

    def setAgent(self, agent:Agent):
        self.agent = agent
    
    def setCodeAgent(self, code_agent:CodeAgent):
        self.code_agent = code_agent
    
    def do_code(self, arg):
        """切换到代码模式。用法: code"""
        if self.code_agent is None:
            print("代码助手未初始化，无法切换到代码模式")
            return
        
        self.current_mode = "code"
        self.prompt = "代码>>> "
        print(f"已切换到代码模式，当前助手: {self.code_agent.name()}")
        print("可用命令:")
        print("  normal - 切换回普通模式")
        print("  analyze <文件路径> - 分析代码文件")
        print("  execute <文件路径或代码> - 执行代码")
        print("  validate <文件路径或代码> - 验证代码")
        print("  help - 查看帮助")
    
    def do_normal(self, arg):
        """切换回普通模式。用法: normal"""
        self.current_mode = "normal"
        self.prompt = ">>> "
        print("已切换回普通模式")

    def _init_history(self):
        """初始化历史记录功能"""
        try:
            if os.path.exists(self.history_file):
                with open(self.history_file, 'r', encoding='utf-8') as f:
                    self.history = f.readlines()
            else:
                os.makedirs(os.path.dirname(self.history_file), exist_ok=True)
        except Exception as e:
            print(f"初始化历史记录失败: {e}") 

    def _save_history(self):
        """保存历史记录到文件"""
        try:
            with open(self.history_file, 'w', encoding='utf-8') as f:
                f.writelines(self.history)
        except Exception as e:
            print(f"保存历史记录失败: {e}")

    def _print_welcome_message(self):
        """打印无外框的大号 AITA Logo 欢迎信息"""
        logo = r"""
          █████╗ ██╗████████╗ █████╗  
         ██╔══██╗██║╚══██╔══╝██╔══██╗ 
         ███████║██║   ██║   ███████║ 
         ██╔══██║██║   ██║   ██╔══██║ 
         ██║  ██║██║   ██║   ██║  ██║ 
         ╚═╝  ╚═╝╚═╝   ╚═╝   ╚═╝  ╚═╝ 
        """

        current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        work_dir = os.getcwd()  # 使用当前工作目录，兼容所有系统

        # 打印大号 Logo 和欢迎信息
        print("\033[1;36m")  # 亮青色
        print(logo)
        print("\033[1m欢迎使用 AITA 助手！\033[0m".center(50))
        print(f"当前时间: {current_time}".center(50))
        print(f"工作目录: {work_dir}".center(50))
        print("\033[0m")  # 重置颜色

    def do_exit(self, arg):
        """退出Shell。用法: exit"""
        print("再见！")
        self._save_history()
        return True  

    def do_quit(self, arg):
        """退出Shell。用法: quit"""
        return self.do_exit(arg)
    
    def do_analyze(self, arg):
        """分析代码文件。用法: analyze <文件路径>"""
        if self.current_mode != "code" or self.code_agent is None:
            print("此命令需要在代码模式下使用")
            return
        
        if not arg.strip():
            print("请提供要分析的文件路径")
            return
        
        print(f"{self.code_agent.name()}: 正在分析代码文件...")
        analysis = self.code_agent.analyze_file(arg.strip())
        
        if analysis["success"]:
            print("分析结果:")
            print(f"  文件: {analysis['file_path']}")
            print(f"  行数: {analysis['lines']}")
            print(f"  字符数: {analysis['size']}")
            print(f"  语言: {analysis['language']}")
            
            if analysis['imports']:
                print(f"  导入: {', '.join(analysis['imports'])}")
            
            if analysis['functions']:
                print(f"  函数: {', '.join(analysis['functions'])}")
            
            if analysis['classes']:
                print(f"  类: {', '.join(analysis['classes'])}")
        else:
            print(f"分析失败: {analysis['error']}")
    
    def do_execute(self, arg):
        """执行代码文件或代码片段。用法: execute <文件路径或代码>"""
        if self.current_mode != "code" or self.code_agent is None:
            print("此命令需要在代码模式下使用")
            return
        
        if not arg.strip():
            print("请提供要执行的文件路径或代码片段")
            return
        
        input_text = arg.strip()
        
        # 判断是文件路径还是代码片段
        if os.path.exists(input_text):
            print(f"{self.code_agent.name()}: 正在执行文件...")
            user_input = f"请执行这个代码文件: {input_text}"
        else:
            print(f"{self.code_agent.name()}: 正在执行代码片段...")
            user_input = f"请执行这段代码:\n```\n{input_text}\n```"
        
        try:
            response = self.code_agent.chat(user_input)
            print(f"{self.code_agent.name()}: {response}")
        except Exception as e:
            print(f"执行失败: {str(e)}")
    
    def do_validate(self, arg):
        """验证代码。用法: validate <文件路径或代码>"""
        if self.current_mode != "code" or self.code_agent is None:
            print("此命令需要在代码模式下使用")
            return
        
        if not arg.strip():
            print("请提供要验证的文件路径或代码片段")
            return
        
        input_text = arg.strip()
        
        # 判断是文件路径还是代码片段
        if os.path.exists(input_text):
            print(f"{self.code_agent.name()}: 正在验证文件...")
            user_input = f"请验证这个代码文件: {input_text}"
        else:
            print(f"{self.code_agent.name()}: 正在验证代码片段...")
            user_input = f"请验证这段代码:\n```\n{input_text}\n```"
        
        try:
            response = self.code_agent.chat(user_input)
            print(f"{self.code_agent.name()}: {response}")
        except Exception as e:
            print(f"验证失败: {str(e)}")
    
    def do_format(self, arg):
        """格式化代码。用法: format <文件路径>"""
        if self.current_mode != "code" or self.code_agent is None:
            print("此命令需要在代码模式下使用")
            return
        
        if not arg.strip():
            print("请提供要格式化的文件路径")
            return
        
        file_path = arg.strip()
        print(f"{self.code_agent.name()}: 正在格式化文件...")
        user_input = f"请格式化这个代码文件: {file_path}"
        
        try:
            response = self.code_agent.chat(user_input)
            print(f"{self.code_agent.name()}: {response}")
        except Exception as e:
            print(f"格式化失败: {str(e)}")  

    def postcmd(self, stop, line):
        """每次命令执行后自动保存历史记录"""
        self._save_history()
        return stop

    def emptyline(self):
        """忽略空行输入"""
        pass

    def default(self, line):
        if self.current_mode == "code" and self.code_agent:
            name = self.code_agent.name()
            if line.lower() == "reset":
                self.code_agent.chat("", reset = True)
                print(f"{name}: 对话历史已重置")
            else:
                print(f"{name}: 正在思考中...", end="", flush=True)
                try:
                    response = self.code_agent.chat(line)
                    print(f"{name}: {response}")
                except Exception as e:
                    print(f"\n错误: {str(e)}")
        elif self.agent:
            name = self.agent.name();
            if line.lower() == "reset":
                self.agent.chat("", reset = True)
                print(f"{name}: 对话历史已重置")
            else:
                print(f"{name}: 正在思考中...", end="", flush=True)
                for chunk in self.agent.chat(line):
                    print(chunk, end="", flush=True)
                print()
        else:
            print("助手未初始化")
